// Tencent is pleased to support the open source community by making
// 蓝鲸智云 - 监控平台/日志平台 (BlueKing - Monitor/Log) available.
// Copyright (C) 2017-2022 THL A29 Limited, a Tencent company. All rights reserved.
// Licensed under the MIT License (the "License"); you may not use this file except in compliance with the License.
// You may obtain a copy of the License at http://opensource.org/licenses/MIT
// Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
// an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.
//

package receiver

import (
	"context"
	"sync"
	"time"

	"github.com/TencentBlueKing/bkmonitor-kits/logger"

	"github.com/TencentBlueKing/bk-collector/confengine"
	"github.com/TencentBlueKing/bk-collector/define"
	"github.com/TencentBlueKing/bk-collector/internal/testutils"
)

type BenchmarkServer struct {
	ctx              context.Context
	cancel           context.CancelFunc
	wg               sync.WaitGroup
	config           Config
	tracesGenerator  *testutils.TracesGenerator
	metricsGenerator *testutils.MetricsGenerator
}

func NewBenchmarkServer(c Config) *BenchmarkServer {
	ctx, cancel := context.WithCancel(context.Background())
	return &BenchmarkServer{
		ctx:              ctx,
		cancel:           cancel,
		config:           c,
		tracesGenerator:  testutils.NewTracesGenerator(c.BenchmarkServer.Traces),
		metricsGenerator: testutils.NewMetricsGenerator(c.BenchmarkServer.Metrics),
	}
}

func (b *BenchmarkServer) Start() error {
	logger.Info("benchmark server start working...")
	conf := b.config.BenchmarkServer
	if conf.Traces.Enabled {
		b.StartTracesBenchmark()
	}
	if conf.Metrics.Enabled {
		b.StartMetricsBenchmark()
	}
	return nil
}

func (b *BenchmarkServer) StartTracesBenchmark() {
	b.wg.Add(1)
	defer b.wg.Done()

	logger.Infof("start to traces benchmarking...")
	ch := make(chan int, 1)
	now := time.Now()
	nums := define.CoreNums()
	for i := 0; i < nums; i++ {
		go b.startTracesBenchmark(ch)
	}

	var total int
	for i := 0; i < nums; i++ {
		total += <-ch
	}
	logger.Infof("traces benchmark result, concurrency: %d, total iteration: %d, take: %v", nums, total, time.Since(now))
}

func (b *BenchmarkServer) startTracesBenchmark(iter chan int) {
	ch := b.tracesGenerator.Ch()
	var n int
	for {
		select {
		case data, ok := <-ch:
			if !ok {
				iter <- n
				return
			}
			PublishRecord(&define.Record{
				RequestType:   define.RequestBenchmark,
				RequestClient: define.RequestClient{IP: "127.0.0.1"},
				RecordType:    define.RecordTraces,
				Data:          data,
			})
			n++

		case <-b.ctx.Done():
			iter <- n
			return
		}
	}
}

func (b *BenchmarkServer) StartMetricsBenchmark() {
	b.wg.Add(1)
	defer b.wg.Done()

	logger.Infof("start to metrics benchmarking...")
	ch := make(chan int, 1)
	now := time.Now()
	nums := define.CoreNums()
	for i := 0; i < nums; i++ {
		go b.startMetricsBenchmark(ch)
	}

	var total int
	for i := 0; i < nums; i++ {
		total += <-ch
	}
	logger.Infof("metrics benchmark result, concurrency: %d, total iteration: %d, take: %v", nums, total, time.Since(now))
}

func (b *BenchmarkServer) startMetricsBenchmark(iter chan int) {
	ch := b.metricsGenerator.Ch()
	var n int
	for {
		select {
		case data, ok := <-ch:
			if !ok {
				iter <- n
				return
			}
			PublishRecord(&define.Record{
				RequestType:   define.RequestBenchmark,
				RequestClient: define.RequestClient{IP: "127.0.0.1"},
				RecordType:    define.RecordMetrics,
				Data:          data,
			})
			n++

		case <-b.ctx.Done():
			iter <- n
			return
		}
	}
}

func (b *BenchmarkServer) Stop() error {
	b.cancel()
	b.wg.Wait()
	return nil
}

func (b *BenchmarkServer) Reload(conf *confengine.Config) error {
	return nil
}
